import os

from openai import models
from paddleocr import PaddleOCR
from paddleocr import PPStructureV3

BASE_PATH = os.getcwd().replace('yl/paddleocr/recognize', '')
REC_MODEL_PATH = BASE_PATH + "models/recognize/"
RESOURCE_PATH = BASE_PATH + "resources/"
OCRv5_RESOURCE_PATH = RESOURCE_PATH + "input/OCRv5/"
STRUCTUREV3_RESOURCE_PATH = RESOURCE_PATH + "input/StructureV3/"

OUTPUT_OCRv5_PATH = RESOURCE_PATH + "output/OCRv5/"
OUTPUT_STRUCTUREv3_PATH = RESOURCE_PATH + "output/StructureV3/"


# OCRv5的识别方法
def img_ocr(img_name_list, orientation_classify_flag=True, doc_unwarping_flag=True, textline_orientation_flag=True):
    ocr = PaddleOCR(
        use_doc_unwarping=doc_unwarping_flag,
        use_doc_orientation_classify=orientation_classify_flag,
        use_textline_orientation=textline_orientation_flag,
    )

    result = ocr.predict(img_name_list,
                         use_doc_unwarping=doc_unwarping_flag,
                         use_doc_orientation_classify=orientation_classify_flag,
                         use_textline_orientation=textline_orientation_flag, )
    suffix = ''
    if orientation_classify_flag is False or doc_unwarping_flag is False or textline_orientation_flag is False:
        suffix = '_不含矫正模型'

    for res in result:
        res.save_to_img(OUTPUT_OCRv5_PATH + "output" + suffix)
        res.save_to_json(OUTPUT_OCRv5_PATH + "output" + suffix)


# StructureV3的识别方法
def structurev3_ocr(img_name_list, seal_recognition_flag=True, table_recognition_flag=True):
    pipeline = PPStructureV3(
        use_seal_recognition=seal_recognition_flag,
        use_table_recognition=table_recognition_flag,
    )
    output = pipeline.predict(img_name_list,
                              use_seal_recognition=seal_recognition_flag,
                              use_table_recognition=table_recognition_flag,
                              )

    suffix = ''
    if seal_recognition_flag is False or table_recognition_flag is False:
        suffix = '_不含矫正模型'

    for res in output:
        res.save_to_html(save_path=OUTPUT_STRUCTUREv3_PATH + "output_" + suffix)
        res.save_to_markdown(save_path=OUTPUT_STRUCTUREv3_PATH + "output_" + suffix)


if __name__ == "__main__":
    # img_ocr(img_name_list=[
    #      OCRv5_RESOURCE_PATH + 'OCRv5测试-标准图片.png',
    #      OCRv5_RESOURCE_PATH + 'OCRv5测试-官方图像矫正图片.png',
    #      OCRv5_RESOURCE_PATH + 'OCRv5测试-官方文本行方向图片.png',
    #      OCRv5_RESOURCE_PATH + 'OCRv5测试-旋转方向图片.png'
    # ], orientation_classify_flag=False,textline_orientation_flag=False)
    structurev3_ocr(img_name_list=[
        # STRUCTUREV3_RESOURCE_PATH + "v3官方版面图片.png",
        STRUCTUREV3_RESOURCE_PATH + "v3官方表格图片.png",
        # STRUCTUREV3_RESOURCE_PATH + "v3官方印章图片.png"
    ])
